// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/string.h>
#include <melon/init/init.h>
#include <gtest/gtest.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <pollux/plan/reference_query_runner.h>
#include <pollux/testing/expression/fuzzer_runner.h>
#include <pollux/testing/expression/spark_special_form_signature_generator.h>
#include <pollux/functions/prestosql/fuzzer/FloorAndRoundArgTypesGenerator.h>
#include <pollux/functions/sparksql/fuzzer/AddSubtractArgTypesGenerator.h>
#include <pollux/functions/sparksql/fuzzer/DivideArgTypesGenerator.h>
#include <pollux/functions/sparksql/fuzzer/MakeTimestampArgTypesGenerator.h>
#include <pollux/functions/sparksql/fuzzer/MultiplyArgTypesGenerator.h>
#include <pollux/functions/sparksql/fuzzer/UnscaledValueArgTypesGenerator.h>
#include <pollux/functions/sparksql/registration/register.h>

using namespace kumo::pollux::functions::sparksql::fuzzer;
using kumo::pollux::fuzzer::ArgTypesGenerator;
using namespace kumo::pollux::plan;

TURBO_FLAG(int64_t,
    seed,
    123456,
    "Initial seed for random number generator "
    "(use it to reproduce previous results).");

using kumo::pollux::fuzzer::FuzzerRunner;

int main(int argc, char** argv) {
  kumo::pollux::functions::sparksql::registerFunctions("");

  ::testing::InitGoogleTest(&argc, argv);

  // Calls common init functions in the necessary order, initializing
  // singletons, installing proper signal handlers for better debugging
  // experience
  melon::Init init(&argc, &argv);

  // The following list are the Spark UDFs that hit issues
  // For rlike you need the following combo in the only list:
  // rlike, md5 and upper
  std::unordered_set<std::string> skipFunctions = {
      "regexp_extract",
      // https://github.com/facebookincubator/pollux/issues/8438
      "regexp_replace",
      "rlike",
      "chr",
      "replace",
      "might_contain",
      "unix_timestamp",
      // from_unixtime throws PolluxRuntimeError when the timestamp is out of the
      // supported range.
      "from_unixtime",
      // timestamp_millis(bigint) can generate timestamps out of the supported
      // range that make other functions throw PolluxRuntimeErrors.
      "timestamp_millis(bigint) -> timestamp",
  };

  // Required by spark_partition_id function.
  std::unordered_map<std::string, std::string> queryConfigs = {
      {kumo::pollux::core::QueryConfig::kSparkPartitionId, "123"},
      {kumo::pollux::core::QueryConfig::kSessionTimezone,
       "America/Los_Angeles"}};

  std::unordered_map<std::string, std::shared_ptr<ArgTypesGenerator>>
      argTypesGenerators = {
          {"add", std::make_shared<AddSubtractArgTypesGenerator>(true)},
          {"add_deny_precision_loss",
           std::make_shared<AddSubtractArgTypesGenerator>(false)},
          {"subtract", std::make_shared<AddSubtractArgTypesGenerator>(true)},
          {"subtract_deny_precision_loss",
           std::make_shared<AddSubtractArgTypesGenerator>(false)},
          {"multiply", std::make_shared<MultiplyArgTypesGenerator>(true)},
          {"multiply_deny_precision_loss",
           std::make_shared<MultiplyArgTypesGenerator>(false)},
          {"divide", std::make_shared<DivideArgTypesGenerator>(true)},
          {"divide_deny_precision_loss",
           std::make_shared<DivideArgTypesGenerator>(false)},
          {"ceil",
           std::make_shared<
               kumo::pollux::exec::test::FloorAndRoundArgTypesGenerator>()},
          {"floor",
           std::make_shared<
               kumo::pollux::exec::test::FloorAndRoundArgTypesGenerator>()},
          {"unscaled_value",
           std::make_shared<UnscaledValueArgTypesGenerator>()},
          {"make_timestamp",
           std::make_shared<MakeTimestampArgTypesGenerator>()}};

  std::shared_ptr<ReferenceQueryRunner> referenceQueryRunner{nullptr};
  return FuzzerRunner::run(
      FLAGS_seed,
      skipFunctions,
      {{}},
      queryConfigs,
      argTypesGenerators,
      {{}},
      referenceQueryRunner,
      std::make_shared<
          kumo::pollux::fuzzer::SparkSpecialFormSignatureGenerator>());
}
